from BeautifulSoup import BeautifulSoup
import StringIO
import Image, ImageDraw
import urllib
import logging

class EdgeBuilder(object):
    """Build title->edge map from a beautifulsoup object
    """
    def __init__(self):
        pass
    
    def split_title(self,str):
        """Split the title by comma
        >>> e=EdgeBuilder()
        >>> e.split_title("K00889 (E2.7.1.68, PIP5K), K01099, R03469, R04404")
        ['K00889', 'K01099', 'R03469', 'R04404']
        """
        def handle_title(title,result):
            t = "".join(title).strip()
            #remove the brackets
            pos = t.find("(")
            if pos!=-1:
                t = t[:pos-1]
            #process the special value K03060-K03048
            if t == "K03060-K03048":
                result.extend(t.split("-"))
                return
            result.append(t)
        result = [] #Return a list of titles
        bracket_count=0
        title = []
        for i in str:
            if i==',' and bracket_count == 0:
                handle_title(title,result)
                title = []
                continue
            elif i=='(':
                bracket_count+=1
            elif i==')':
                bracket_count-=1    
            title.append(i)
        handle_title(title,result)
        return result
            
    def process_html(self,edges):
        """
        Return title->edges map and edges set
        soup: is the soup object generated from the orgianl html source code
        """
        print "Processing html"
        title = {}  #title to Edges map. eg. K00611->edges object
        edge_set = set()
        n_edges = 0
        for i in edges:
            n_edges +=1
            #Get all title for this edge
            titles = self.split_title(i['title'])
            #Get the polygon 
            poly = i['coords']
            #Create Edge object
            edge = Edge(poly,titles,i)
            edge_set.add(edge)
            #Insert edge to title->edges map
            for t in titles:
                title.setdefault(t,set()).add(edge)
        #print n_edges,"processed"
        return title, edge_set
                
class Edge(object):
    def __init__(self,poly_path,title,soup_node):
        self.poly = [int(i) for i in poly_path.split(',')]
        self.soup_node = soup_node
        self.title = title
        assert(len(title)>=1 and len(self.poly)>=4)
        self.hit_read = 0
        self.specie = set()
        self._color = 0.0
        self.n_kval = 0

    def clear(self):
        self._color = 0.0
        self.n_kval = 0
        self.hit_read = 0
        self.specie = set()

    def update_avg(self,distance):
        self._color = \
                    (self._color*self.n_kval + distance)/(self.n_kval+1)
        self.n_kval+=1

    def update_max(self, distance):
        if self._color<distance:
            self._color = distance

    def diversity(self,):
        return self._color
    
    def to_html(self,value):
        if value>0:
            orig_tag = str(self.soup_node)
            orig_arr = orig_tag.split("title=\"")
            assert(len(orig_arr)==2)
            if type(value) == int:
                read = "title=\"%d "%value
            else:
                read = "title=\"%f "%value
            return orig_arr[0]+read+orig_arr[1]
        return str(self.soup_node)
        
    def number_of_hit(self):
        return self.hit_read

    def add_hit(self,read):
        self.hit_read += 1

    def add_spec(self,spec):
        self.specie.add(spec)

    def number_of_spec(self):
        return len(self.specie)

    def __repr__(self):
        return ",".join(self.title)+":"+str(self.hit_read)
        

def read_div_file(file):
    "Read spediv file and return a {K value-> diversity} map"
    k_val_map = {}
    
    for line in file:
        if not line.startswith("#"):
            fields = line.split()
            if len(fields)!=4:
                logging.warning(file.name+":"+line+" Does not have 4 fields,Skip")
                continue
            k_val = fields[0]
            size = float(fields[3])
            k_val_map[k_val]=size
    return k_val_map

def mark_edge_by_div(file, title):
    "Read file and mark edge by diversity"
    marked = set()
    max_dis = 0.0
    k_val_map = read_div_file(file)
    for k in k_val_map:
        t = title.get(k,[])
        for edge in t:
            #if k_val_map[k] > max_dis:
            #    max_dis = k_val_map[k]
            #edge.update_max(k_val_map[k])
            edge.update_avg(k_val_map[k])
            if edge.diversity() > max_dis:
                max_dis = edge.diversity()
            marked.add(edge)
    return marked, "diversity", max_dis

def mark_edge(read_lst, edge_map):
    "Read file and Mark edge by reads"
    mark_set = set()
    max_read = 0
    for line in read_lst:
        if "#" in line :
            continue
        split_result = line.strip().split()
        if len(split_result)<2:
            continue
        read,ks = split_result[0],split_result[1]
        if "+" in ks:
            ks = line.split("+")
        else:
            ks = [ks]
        #assert(edge_map.has_key(i))
        for i in ks:
            if not edge_map.has_key(i):
                continue
            edge_set = edge_map[i]
            for every_edge in edge_set:
                every_edge.add_hit(read)
                if every_edge.number_of_hit()>max_read:
                    max_read = every_edge.number_of_hit()
            mark_set = mark_set.union(edge_set)
    return mark_set,"number_of_hit",float(max_read)
    

def mark_edge_spec(read_lst, edge_map):
    "Read file and mark edge by species count"
    mark_set = set()
    max_read = 0
    for line in read_lst:
        if "#" in line :
            continue
        split_result = line.strip().split()
        if len(split_result) != 3:
            continue
        read,i,spec = split_result[0],split_result[1],split_result[2]
        #assert(edge_map.has_key(i))
        if not edge_map.has_key(i):
            continue
        edge_set = edge_map[i]
        for every_edge in edge_set:
            every_edge.add_spec(spec)
            if every_edge.number_of_spec()>max_read:
                max_read = every_edge.number_of_spec()
        mark_set = mark_set.union(edge_set)
    return mark_set,"number_of_spec",float(max_read)

def draw_image(im,marked,method,max_val):
    """
    Draw image using a method in Edge object
    im: Image
    marked: Marked edge
    method: a string for method name
    max_val: max value for makred edge
    """
    draw = ImageDraw.Draw(im)
    for i in marked:
        rgb = val_to_rgb(getattr(i,method)()/max_val*100)
        draw.polygon(i.poly,fill=rgb)
    del draw
    image_file = StringIO.StringIO()
    im.save(image_file,"PNG")
    image_src = 'data:image/png,'+urllib.quote(image_file.getvalue())
    poly_arr = []
    for edge in g_edge_set:
        poly_arr.append(edge.to_html(getattr(edge,method)()))
        edge.clear()
    poly_html = "\n".join(poly_arr)
    return image_src, poly_html

def draw_image_proportion(im,marked,method,max_val):
    """
    Draw image using a method in Edge object
    Color the edge by proportion
    im: Image
    marked: Marked edge
    method: a string for method name
    max_val: max value for makred edge
    """
    draw = ImageDraw.Draw(im)
    n_edge = len(marked)
    n_each_color = n_edge/10 - 1
    color = 0
    color_list = []
    value_list = []
    counter = 0
    for counter,i in enumerate(sorted(marked,key=lambda x:getattr(x,method)())):
        rgb = val_to_rgb(color)
        draw.polygon(i.poly,fill=rgb)
        if counter%n_each_color==0 and color < 100:
            color_list.append(rgb)
            value_list.append(getattr(i,method)())
            color += 10
    del draw
    color_list.append(rgb)
    value_list.append(max_val)
    color_bar(im,color_list,value_list,method)
    image_file = StringIO.StringIO()
    im.save(image_file,"PNG")
    image_src = 'data:image/png,'+urllib.quote(image_file.getvalue())
    poly_arr = []
    for edge in g_edge_set:
        poly_arr.append(edge.to_html(getattr(edge,method)()))
        edge.clear()
    poly_html = "\n".join(poly_arr)
    return image_src, poly_html

#def color_bar(im,color,val, title):
#    assert(len(color)==len(val)),(color,val)
#    start=(60,10)
#    end=(start[0]+220,start[1]+130)    
#    draw = ImageDraw.Draw(im)
#    draw.rectangle([start,end],fill="white",outline="black")
#    draw.text(start,"     Color             "+title,fill="black")
#    last_num = 0.0
#    for i,c in enumerate(zip(color,val)):
#        if type(c[1])==float:
#            fmt = "%0.3f - %0.3f"%(last_num,c[1])
#        else:
#            fmt = "%d - %d"%(last_num,c[1])
#        draw.line([start[0]+10, start[1]+10*i+20, end[0]-100, start[1]+10*i+20],fill=c[0])
#        draw.text((end[0]-90,start[1]+10*i+15),fmt,fill="black")
#        last_num = c[1]
#    del draw

def color_bar(im,color,val, title):
    assert(len(color)==len(val)),(color,val)
    start=(1300,10)
    end=(start[0]+530,start[1]+35)    
    draw = ImageDraw.Draw(im)
    draw.rectangle([start,end],fill="white",outline="black")
    first_line = (start[0]+10,start[1]+10)
    draw.text(first_line,"Color",fill="black")
    second_line = (start[0]+1,first_line[1]+10)
    draw.text(second_line,title,fill="black")
    last_x = start[0]+80
    if type(val[0])==int:
        fmt = "%d"
    else:
        fmt = "%0.2f" 
    for i,c in enumerate(zip(color,val)):
        draw.line([last_x, start[1]+25, last_x+40, start[1]+25],fill=c[0],width = 5)
        draw.text((last_x+20,start[1]+10),fmt%c[1],fill="black")
        last_x += 40
    del draw

    
def val_to_rgb(val):
    """Return the RGB value
    val: range 0-100 (blue to red)
    """
    assert(0<=val<=100),val
    if val<25:
        R=0
        G=val/25.0*255
        B=255
    elif val<50:
        R=0
        G=255
        B=(50-val)/25.0*255
    elif val<75:
        R=(val-50)/25.0*255
        G=255
        B=0
    else:
        R=255
        G=(100-val)/25.0*255
        B=0
    return "rgb(%d,%d,%d)"%(R,G,B)

path = open("../show_pathway.htm",'r')
soup = BeautifulSoup(path)
builder = EdgeBuilder()
polys = soup.findAll(shape='poly')
title_map,g_edge_set = builder.process_html(polys)
    


